# Get version of CUDA and enable it for compilation if CUDA > 11.0
# This solves https://github.com/IDEA-Research/Grounded-Segment-Anything/issues/53
# and https://github.com/IDEA-Research/Grounded-Segment-Anything/issues/84
# when running in Docker
# Check if nvcc is installed
NVCC := $(shell which nvcc)
BASE_IMAGE := "winsleo/vnc-desktop:pytorch2.1.2-cuda12.1-cudnn8-devel"
SAM_DATA_DIR ?= /workspace/dataset

ifeq ($(NVCC),)
	# NVCC not found
	USE_CUDA := 0
	NVCC_VERSION := "not installed"
else
	NVCC_VERSION := $(shell nvcc --version | grep -oP 'release \K[0-9.]+')
	USE_CUDA := $(shell echo "$(NVCC_VERSION) > 11" | bc -l)
endif

# Add the list of supported ARCHs
ifeq ($(USE_CUDA), 1)
	TORCH_CUDA_ARCH_LIST := "5.0;6.0;6.1;7.0;7.5;8.0;8.6;8.7;8.9+PTX"
	BUILD_MESSAGE := "I will try to build the image with CUDA support"
else
	TORCH_CUDA_ARCH_LIST :=
	BUILD_MESSAGE := "CUDA $(NVCC_VERSION) is not supported"
endif

build-image:
	@echo $(BUILD_MESSAGE)
	docker build --build-arg USE_CUDA=$(USE_CUDA) \
	--build-arg TORCH_ARCH=$(TORCH_CUDA_ARCH_LIST) \
	--build-arg BASE_IMAGE=$(BASE_IMAGE) \
	-t gsa:v0 .

run:
ifeq (,$(wildcard ./sam_vit_h_4b8939.pth))
	wget --no-check-certificate https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
endif
ifeq (,$(wildcard ./sam_hq_vit_h.pth))
	wget --no-check-certificate https://huggingface.co/lkeab/hq-sam/resolve/main/sam_hq_vit_h.pth
endif
ifeq (,$(wildcard ./groundingdino_swint_ogc.pth))
	wget --no-check-certificate https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth
endif
ifeq (,$(wildcard ./bert-base-uncased))
	mkdir bert-base-uncased && cd bert-base-uncased && \
	wget --no-check-certificate https://huggingface.co/google-bert/bert-base-uncased/resolve/main/config.json && \
	wget --no-check-certificate https://huggingface.co/google-bert/bert-base-uncased/resolve/main/pytorch_model.bin && \
	wget --no-check-certificate https://huggingface.co/google-bert/bert-base-uncased/resolve/main/tokenizer_config.json && \
	wget --no-check-certificate https://huggingface.co/google-bert/bert-base-uncased/resolve/main/tokenizer.json && \
	wget --no-check-certificate https://huggingface.co/google-bert/bert-base-uncased/resolve/main/vocab.txt
endif

	docker run --name gsa --privileged\
		-it --rm --gpus all --detach\
		--network=bridge \
		--expose=22 \
		--expose=5900 \
		--expose=6080 \
		--publish=10022:22 \
		--publish=5901:5900 \
		--publish=6081:6080 \
		-v /tmp/.X11-unix:/tmp/.X11-unix \
		-v "${PWD}":/workspace/Grounded-Segment-Anything \
		-v "${SAM_DATA_DIR}":/workspace/dataset \
		-e DISPLAY=$DISPLAY \
		gsa:v0

clean:
	docker rm -f gsa